package org.easy.ssh;

import cn.hutool.core.thread.ExecutorBuilder;
import com.alibaba.fastjson.JSON;
import com.jcraft.jsch.Channel;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.websocket.Constant;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Map;
import java.util.Objects;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * @Description:
 * @Author: wxl
 * @CreateTime: 2023/2/16 8:15
 */
public class EasyWebSshOperator implements IEasyWebSshOperator{

    private static final Map<String, SshConnectInfo> SSH_CONFIG_MAP = new ConcurrentHashMap<>();

    private static final Logger LOGGER = LoggerFactory.getLogger(EasyWebSshOperator.class);

    private final ThreadPoolExecutor executor = ExecutorBuilder.create().build();

    public EasyWebSshOperator(){}

    @Override
    public void initConnection(WebSocketSession session) {
        JSch jSch = new JSch();
        SshConnectInfo info = new SshConnectInfo(jSch,session);
        String uuid = String.valueOf(session.getAttributes().get(Constant.USER_UUID_KEY));
        SSH_CONFIG_MAP.put(uuid,info);
    }

    @Override
    public void recvHandler(String data, WebSocketSession session) {
        WebSshData webSshData = JSON.parseObject(data,WebSshData.class);
        String uuid = String.valueOf(session.getAttributes().get(Constant.USER_UUID_KEY));
        SshConnectInfo info = SSH_CONFIG_MAP.get(uuid);
        if (Constant.WEB_SSH_OPERATE_COMMAND.equals(webSshData.getOperate())) {
            String command = webSshData.getCommand();
            if (Objects.nonNull(info)) {
                try {
                    transToSsh(info.getChannel(),command);
                } catch (IOException e) {
                    LOGGER.error("发送命令异常:{}",e.getMessage());
                }
            }
        } else if (Constant.WEB_SSH_OPERATE_CONNECT.equals(webSshData.getOperate())) {
            final WebSshData finalWebSshData = webSshData;
            executor.execute(() -> {
                try {
                    connectToSsh(info,finalWebSshData,session);
                }catch (JSchException | IOException e) {
                    LOGGER.error("连接SSH异常:{}",e.getMessage());
                }
            });
        }
    }

    private void connectToSsh(SshConnectInfo info, WebSshData finalWebSshData, WebSocketSession session) throws JSchException, IOException {
        Properties config = new Properties();
        config.put("StrictHostKeyChecking","no");
        Session sshSession = info.getJSch().getSession(finalWebSshData.getUsername(),finalWebSshData.getHost(),finalWebSshData.getPort());
        sshSession.setConfig(config);
        sshSession.setPassword(finalWebSshData.getPassword());
        sshSession.connect(30000);

        final Channel channel = sshSession.openChannel("shell");
        channel.connect(3000);
        info.setChannel(channel);

        transToSsh(channel,"\r");

        try (InputStream outStream = channel.getInputStream()) {
            //循环读取
            byte[] buffer = new byte[1024];
            int i;
            //如果没有数据来，线程会一直阻塞在这个地方等待数据。
            while ((i = outStream.read(buffer)) != -1) {
                sendMessage(Arrays.copyOfRange(buffer, 0, i), session);
            }
        } finally {
            sshSession.disconnect();
            channel.disconnect();
        }

    }

    @Override
    public void sendMessage(byte[] buffer, WebSocketSession session) throws IOException {
        session.sendMessage(new TextMessage(buffer));
    }

    @Override
    public void close(WebSocketSession session) {
        String uuid = String.valueOf(session.getAttributes().get(Constant.USER_UUID_KEY));
        SshConnectInfo info = SSH_CONFIG_MAP.get(uuid);
        if (Objects.nonNull(info)) {
            if (Objects.nonNull(info.getChannel())) {
                info.getChannel().disconnect();
            }
            SSH_CONFIG_MAP.remove(uuid);
        }
    }

    /**
     * 发送命令
     * @param channel 通道
     * @param command 命令
     * @throws IOException 异常
     */
    private void transToSsh(Channel channel, String command) throws IOException {
        if (Objects.nonNull(channel)) {
            OutputStream inStream = channel.getOutputStream();
            inStream.write(command.getBytes(StandardCharsets.UTF_8));
            inStream.flush();
        }
    }

}
